# -*- coding: utf-8 -*-
import numpy as np

#%%
out_file = "../hls/skynet_para_256.h"
bitw, sym_bitw = 5,  "BIT_WEIT"
bitb, sym_bitb = 16, "BIT_BIAS"
bitm, sym_bitm = 16, "BIT_MULT"

#%%
def int2hex(x):
    return '\"{:#x}\"'.format(x)

def write_const(file, symbol, val):
    code = "constexpr unsigned {} = {:d};\n".format(symbol, val)
    file.write(code)
    return

def write_line(file, data, bit):
    assert len(data.shape) == 2
    data[data < 0] += np.power(2, bit)
    outer, inner = data.shape
    line = [None] * outer
    file.write("\t")
    for i in range(outer):
        cat = 0
        for j in range(inner):
            cat |= int(data[i, j]) << (bit * j)
        line[i] = cat
    file.write(", ".join(map(int2hex, line)))
    file.write(",\n")
    return

def write_arr1D(file, data, bit, sym_arr, sym_width, sym_dim):
    sp = data.shape
    assert len(sp) == 2
    if sp[0] % 8 == 0:
        data = data.reshape(sp[0] // 8, 8, sp[1])
    else:
        data = data.reshape(1, sp[0], sp[1])
    title = "const ap_uint<{0}> {1}[{2}] = {{\n".format(sym_width, sym_arr, sym_dim)
    file.write(title)
    for i in range(data.shape[0]):
        write_line(file, data[i], bit)
    file.write("};\n")
    return

def write_arr2D(file, data, bit, sym_arr, sym_width, sym_dim1, sym_dim2):
    assert len(data.shape) == 3
    title = "const ap_uint<{0}> {1}[{2}][{3}] = {{\n".format(sym_width, sym_arr, sym_dim1, sym_dim2)
    file.write(title)
    for i in range(data.shape[0]):
        write_line(file, data[i], bit)
    file.write("};\n")
    return

def write_DW_weight(file, weight, bias, mult, n_io, n_act, tag):
    n_ch = weight.shape[2]
    assert len(weight.shape) == 3
    assert weight.shape[0] == weight.shape[1] == 3
    assert len(bias.shape) == len(mult.shape) == 1
    assert bias.shape[0] == mult.shape[0] == n_ch
    assert n_ch >= n_io and n_ch % n_io == 0
    assert n_ch >= n_act and n_ch % n_act == 0
    
    data = weight.reshape(9, n_ch // n_io, n_io).transpose(1, 0, 2)
    assert data.shape == (n_ch // n_io, 9, n_io)
    # data[data < 0] += np.power(2, bitw)
    
    sym_arr = "{}_DW".format(tag)
    sym_nch = "{}_NCH".format(sym_arr)
    sym_nio = "{}_NIO".format(sym_arr)
    sym_act = "{}_NACT".format(sym_arr)
    write_const(file, sym_nch, n_ch)
    write_const(file, sym_nio, n_io)
    write_const(file, sym_act, n_act)
   
    sym_width = "{} * {}".format(sym_nio, sym_bitw)
    sym_dim = "{} / {}".format(sym_nch, sym_nio)
    write_arr2D(file, data, bitw, sym_arr, sym_width, sym_dim, 9)
    
    bias = bias.reshape(n_ch // n_act, n_act)
    sym_arr = "{}_DB".format(tag)
    sym_width = "{} * {}".format(sym_act, sym_bitb)
    sym_dim = "{} / {}".format(sym_nch, sym_act)
    write_arr1D(file, bias, bitb, sym_arr, sym_width, sym_dim)
    
    mult = mult.reshape(n_ch // n_act, n_act)
    sym_arr = "{}_DM".format(tag)
    sym_width = "{} * {}".format(sym_act, sym_bitm)
    sym_dim = "{} / {}".format(sym_nch, sym_act)
    write_arr1D(file, mult, bitm, sym_arr, sym_width, sym_dim)
    
    file.write("\n")
    return

def write_PW_weight(file, weight, bias, mult, n_in, n_out, n_act, tag):
    n_och, n_ich = weight.shape
    assert len(weight.shape) == 2
    assert len(bias.shape) == len(mult.shape) == 1
    assert bias.shape[0] == mult.shape[0] == n_och
    assert n_och >= n_out and n_och % n_out == 0
    assert n_ich >= n_in  and n_ich % n_in  == 0
    if n_act > 0:
        assert n_och >= n_act and n_och % n_act == 0
    
    data = weight.reshape(n_och // n_out, n_out, n_ich // n_in, n_in).transpose(0, 2, 1, 3)
    data = data.reshape(n_och // n_out, n_ich // n_in, n_out * n_in)
    # data[data < 0] += np.power(2, bitw)
    
    sym_arr = "{}_PW".format(tag)
    sym_ich = "{}_NICH".format(sym_arr)
    sym_och = "{}_NOCH".format(sym_arr)
    sym_in  =  "{}_NIN".format(sym_arr)
    sym_out = "{}_NOUT".format(sym_arr)
    sym_act = "{}_NACT".format(sym_arr)
    write_const(file, sym_ich, n_ich)
    write_const(file, sym_och, n_och)
    write_const(file, sym_in , n_in )
    write_const(file, sym_out, n_out)
    if n_act > 0:
        write_const(file, sym_act, n_act)
    
    sym_width = "{} * {} * {}".format(sym_out, sym_in, sym_bitw)
    sym_dim = "{} / {}".format(sym_och, sym_out)
    sym_dim2 = "{} / {}".format(sym_ich, sym_in)
    write_arr2D(file, data, bitw, sym_arr, sym_width, sym_dim, sym_dim2)
    
    if n_act > 0:
        bias = bias.reshape(n_och // n_act, n_act)
        sym_arr = "{}_PB".format(tag)
        sym_width = "{} * {}".format(sym_act, sym_bitb)
        sym_dim = "{} / {}".format(sym_och, sym_act)
        write_arr1D(file, bias, bitb, sym_arr, sym_width, sym_dim)
        
        mult = mult.reshape(n_och // n_act, n_act)
        sym_arr = "{}_PM".format(tag)
        sym_width = "{} * {}".format(sym_act, sym_bitm)
        sym_dim = "{} / {}".format(sym_och, sym_act)
        write_arr1D(file, mult, bitm, sym_arr, sym_width, sym_dim)
    
    file.write("\n")
    return

def write_BM_weight(file, bias, mult, n_io, symbol):
    return

#%%
weight = np.load("../data/skynet_weight_int32_KKC_OCIC.npz")

with open(out_file, "w") as f:
    f.write("#pragma once\n")
    f.write("#ifndef SKYNET_PARA_H_\n")
    f.write("#define SKYNET_PARA_H_\n")
    # f.write("#include <ap_int.h>\n\n")
    f.write("#include \"skynet_flow.h\"\n\n")
    
    # write_const(f, sym_bitw, bitw)
    # write_const(f, sym_bitb, bitb)
    # write_const(f, sym_bitm, bitm)
    # f.write("\n")
    
    # Bundle 1
    write_DW_weight(f, weight["L0-W"], weight["L0-B"], weight["L0-M"], 3, 1, "L0")
    write_PW_weight(f, weight["L1-W"], weight["L1-B"], weight["L1-M"], 3, 4, 4, "L0")
    
    # Bundle 2
    write_DW_weight(f, weight["L2-W"], weight["L2-B"], weight["L2-M"], 12, 1, "L1")
    write_PW_weight(f, weight["L3-W"], weight["L3-B"], weight["L3-M"], 8, 12, 2, "L1")
    
    # Bundle 3
    write_DW_weight(f, weight["L4-W"], weight["L4-B"], weight["L4-M"], 6, 1, "L2")
    write_PW_weight(f, weight["L5-W"], weight["L5-B"], weight["L5-M"], 8, 12, 1, "L2")
    
    # Bundle 4
    write_DW_weight(f, weight["L6-W"], weight["L6-B"], weight["L6-M"], 3, 1, "L3")
    write_PW_weight(f, weight["L7-W"], weight["L7-B"], weight["L7-M"], 8, 12, 1, "L3")
    
    # Bundle 5
    write_DW_weight(f, weight["L8-W"], weight["L8-B"], weight["L8-M"], 6, 1, "L4")
    write_PW_weight(f, weight["L9-W"], weight["L9-B"], weight["L9-M"], 16, 16, 1, "L4")
    
    # Bundle 6
    write_DW_weight(f, weight["L10-W"], weight["L10-B"], weight["L10-M"], 16, 2, "L5")
    write_PW_weight(f, weight["L11-W"], weight["L11-B"], weight["L11-M"], 10, 16, 1, "L5")
    
    # Bundle 7
    write_PW_weight(f, weight["L12-W"], weight["L12-B"], weight["L12-M"], 2, 1, -1, "L6")
    
    f.write("#endif\n")


